package com.cube.algorithms.heuristic;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.stream.Collectors;

/**
 * @author litb
 * @since 2025/6/18 20:50
 * <p>
 * 使用蚁群优化算法 (Ant Colony Optimization, ACO) 解决旅行商问题 (TSP) 的示例。
 */

public class TspAntColonyOptimizationExample {

    // --- ACO 核心参数 ---
    /**
     * alpha: 信息素重要性因子。值越大，蚂蚁越倾向于选择之前走过的路径。
     */
    private final double alpha = 1.0;
    /**
     * beta: 启发信息（距离）重要性因子。值越大，蚂蚁越倾向于选择距离更近的城市。
     */
    private final double beta = 5.0;
    /**
     * rho: 信息素蒸发率。表示信息素随时间挥发的程度。
     */
    private final double evaporationRate = 0.5;
    /**
     * Q: 信息素释放总量。一个常数，用于计算蚂蚁在路径上释放的信息素量。
     */
    private final double pheromoneDepositAmount = 100.0;
    /**
     * 蚂蚁数量
     */
    private final int numAnts;
    /**
     * 迭代次数
     */
    private final int numIterations;


    // --- 算法内部数据结构 ---
    private final List<City> cities;
    private final int numCities;
    private final double[][] pheromones;
    private final double[][] distances;
    private final Random random = new Random();

    private List<Integer> bestTourOrder;
    private double bestTourLength = Double.MAX_VALUE;


    /**
     * 构造函数，初始化算法。
     *
     * @param cities        城市列表
     * @param numAnts       蚂蚁数量
     * @param numIterations 迭代次数
     */
    public TspAntColonyOptimizationExample(List<City> cities, int numAnts, int numIterations) {
        this.cities = cities;
        this.numCities = cities.size();
        this.numAnts = numAnts;
        this.numIterations = numIterations;
        this.distances = new double[numCities][numCities];
        this.pheromones = new double[numCities][numCities];
        setup();
    }

    // --- 主程序入口 ---
    public static void main(String[] args) {
        // 创建城市仓库
        List<City> cities = new ArrayList<>();
        cities.add(new City("北京", 116.46, 39.92));
        cities.add(new City("上海", 121.48, 31.22));
        cities.add(new City("广州", 113.23, 23.16));
        cities.add(new City("深圳", 114.07, 22.62));
        cities.add(new City("成都", 104.06, 30.67));
        cities.add(new City("杭州", 120.19, 30.26));
        cities.add(new City("西安", 108.95, 34.27));
        cities.add(new City("武汉", 114.31, 30.52));
        cities.add(new City("重庆", 106.54, 29.59));
        cities.add(new City("天津", 117.2, 39.13));

        // 创建算法实例并执行
        // 参数: 城市列表, 蚂蚁数量, 迭代次数
        TspAntColonyOptimizationExample aco = new TspAntColonyOptimizationExample(cities, 20, 200);
        aco.solve();

        // 打印最终结果
        List<City> bestTour = aco.getBestTour();
        System.out.println("蚁群算法找到的最优路径:");
        String pathString = bestTour.stream()
                .map(City::getName)
                .collect(Collectors.joining(" -> "));
        System.out.println(pathString + " -> " + bestTour.get(0).getName());
        System.out.printf("总距离: %.2f%n", aco.bestTourLength);
    }

    /**
     * 准备工作：计算城市间距离，初始化信息素矩阵。
     */
    private void setup() {
        // 预先计算所有城市间的距离
        for (int i = 0; i < numCities; i++) {
            for (int j = i; j < numCities; j++) {
                double dist = cities.get(i).distanceTo(cities.get(j));
                distances[i][j] = dist;
                distances[j][i] = dist;
            }
        }

        // 初始化所有路径上的信息素为一个较小的正数
        double initialPheromone = 1.0 / numCities;
        for (int i = 0; i < numCities; i++) {
            Arrays.fill(pheromones[i], initialPheromone);
        }
    }

    /**
     * 运行蚁群算法求解。
     */
    public void solve() {
        // 进行指定次数的迭代
        for (int i = 0; i < numIterations; i++) {
            // 为每只蚂蚁生成一条路径
            List<List<Integer>> antTours = new ArrayList<>();
            for (int ant = 0; ant < numAnts; ant++) {
                antTours.add(buildTour());
            }

            // 更新信息素
            updatePheromones(antTours);

            // 更新全局最优解
            updateBestTour(antTours);

            // （可选）打印当前迭代的最优结果
            // System.out.printf("Iteration %d: Best tour length = %.2f%n", i, bestTourLength);
        }
    }

    /**
     * 为一只蚂蚁构建一条完整的旅行路径。
     *
     * @return 代表路径的城市索引列表
     */
    private List<Integer> buildTour() {
        List<Integer> tour = new ArrayList<>();
        boolean[] visited = new boolean[numCities];
        // 随机选择一个起始城市
        int currentCity = random.nextInt(numCities);
        tour.add(currentCity);
        visited[currentCity] = true;

        // 逐步构建路径，直到访问完所有城市
        while (tour.size() < numCities) {
            currentCity = selectNextCity(currentCity, visited);
            tour.add(currentCity);
            visited[currentCity] = true;
        }
        return tour;
    }

    /**
     * 核心逻辑：根据信息素和启发信息，概率性地选择下一个城市。
     *
     * @param currentCity 当前所在城市
     * @param visited     已访问城市标记数组
     * @return 下一个城市的索引
     */
    private int selectNextCity(int currentCity, boolean[] visited) {
        double[] probabilities = new double[numCities];
        double probabilitiesSum = 0.0;

        // 计算所有未访问邻居的转移概率
        for (int i = 0; i < numCities; i++) {
            if (!visited[i]) {
                // 结合信息素和启发信息（距离倒数）
                double pheromoneLevel = Math.pow(pheromones[currentCity][i], alpha);
                double heuristicLevel = Math.pow(1.0 / distances[currentCity][i], beta);
                probabilities[i] = pheromoneLevel * heuristicLevel;
                probabilitiesSum += probabilities[i];
            }
        }

        // 轮盘赌选择
        double rand = random.nextDouble() * probabilitiesSum;
        double cumulativeProbability = 0.0;
        for (int i = 0; i < numCities; i++) {
            if (!visited[i]) {
                cumulativeProbability += probabilities[i];
                if (cumulativeProbability >= rand) {
                    return i;
                }
            }
        }

        // 容错：如果轮盘赌失败（由于浮点精度问题），则随机选择一个未访问的
        for (int i = 0; i < numCities; i++) {
            if (!visited[i]) {
                return i;
            }
        }
        return -1; // Should not happen
    }

    /**
     * 更新全局信息素矩阵。
     *
     * @param antTours 当前迭代中所有蚂蚁生成的路径
     */
    private void updatePheromones(List<List<Integer>> antTours) {
        // 1. 信息素蒸发
        for (int i = 0; i < numCities; i++) {
            for (int j = 0; j < numCities; j++) {
                pheromones[i][j] *= (1.0 - evaporationRate);
            }
        }

        // 2. 信息素释放
        for (List<Integer> tour : antTours) {
            double tourLength = calculateTourLength(tour);
            double pheromoneToAdd = pheromoneDepositAmount / tourLength;
            for (int i = 0; i < numCities - 1; i++) {
                int from = tour.get(i);
                int to = tour.get(i + 1);
                pheromones[from][to] += pheromoneToAdd;
                // 对称路径
                pheromones[to][from] += pheromoneToAdd;
            }
            // 连接路径终点和起点
            int last = tour.get(numCities - 1);
            int first = tour.get(0);
            pheromones[last][first] += pheromoneToAdd;
            pheromones[first][last] += pheromoneToAdd;
        }
    }

    /**
     * 在当前迭代的所有路径中，寻找最优路径并更新全局最优解。
     */
    private void updateBestTour(List<List<Integer>> antTours) {
        for (List<Integer> tour : antTours) {
            double currentLength = calculateTourLength(tour);
            if (currentLength < bestTourLength) {
                bestTourLength = currentLength;
                bestTourOrder = new ArrayList<>(tour);
            }
        }
    }

    /**
     * 计算一条路径的总长度。
     */
    private double calculateTourLength(List<Integer> tour) {
        double length = 0;
        for (int i = 0; i < numCities - 1; i++) {
            length += distances[tour.get(i)][tour.get(i + 1)];
        }
        length += distances[tour.get(numCities - 1)][tour.get(0)];
        return length;
    }

    /**
     * 获取最终找到的最优路径。
     */
    public List<City> getBestTour() {
        if (bestTourOrder == null) {
            return null;
        }
        return bestTourOrder.stream().map(cities::get).collect(Collectors.toList());
    }

    /**
     * 城市类
     */
    public static class City {
        private final String name;
        private final double x;
        private final double y;

        public City(String name, double x, double y) {
            this.name = name;
            this.x = x;
            this.y = y;
        }

        public String getName() {
            return name;
        }

        public double distanceTo(City other) {
            double deltaX = this.x - other.x;
            double deltaY = this.y - other.y;
            return Math.sqrt(deltaX * deltaX + deltaY * deltaY);
        }

        @Override
        public String toString() {
            return getName();
        }
    }
} 